from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional, List
import uvicorn
import os
import shutil
from datetime import datetime
from face_searcher_milvus import FaceSearcherMilvus as FaceSearcher
import logging
import re
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.openapi.utils import get_openapi
import asyncio
import signal
import sys
import random
import string

# 从环境变量获取环境类型
ENV = os.getenv("APP_ENV", "production")

# 根据环境决定是否启用 swagger
if ENV == "development":
    app = FastAPI(
        title="人脸搜索服务",
        description="提供人脸特征提取、存储和相似度搜索服务",
        version="1.0.0"
    )
else:
    app = FastAPI(
        docs_url=None,    # 禁用 swagger UI
        redoc_url=None,   # 禁用 redoc
        openapi_url=None  # 禁用 openapi schema
    )

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化人脸搜索器
# searcher = FaceSearcher(mongo_uri="mongodb://192.168.0.234:27017/")
searcher = FaceSearcher()

# 添加文件存储相关的配置
TEMP_UPLOAD_DIR = "temp_uploads"  # 临时文件目录
PERMANENT_UPLOAD_DIR = "uploads"   # 永久存储目录
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png'}

# 创建必要的目录
os.makedirs(TEMP_UPLOAD_DIR, exist_ok=True)
os.makedirs(PERMANENT_UPLOAD_DIR, exist_ok=True)

def generate_random_string(length=4):
    """生成指定长度的随机字符串"""
    return ''.join(random.choices(string.ascii_letters + string.digits, k=length))

def save_upload_file(upload_file: UploadFile, permanent: bool = False) -> str:
    """
    保存上传的文件并返回文件路径
    
    Args:
        upload_file: 上传的文件
        permanent: 是否永久保存
        
    Returns:
        str: 保存的文件路径
    """
    # 检查文件扩展名
    ext = os.path.splitext(upload_file.filename)[1].lower()
    if ext not in IMAGE_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件格式。支持的格式: {', '.join(IMAGE_EXTENSIONS)}"
        )
    
    # 生成文件名：时间戳_随机数_原文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    random_str = generate_random_string(4)
    filename = f"{timestamp}_{random_str}_{upload_file.filename}"
    
    # 选择存储目录
    save_dir = PERMANENT_UPLOAD_DIR if permanent else TEMP_UPLOAD_DIR
    file_path = os.path.join(save_dir, filename)
    
    # 保存文件
    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(upload_file.file, buffer)
        return file_path
    except Exception as e:
        logger.error(f"保存文件失败: {str(e)}")
        raise HTTPException(status_code=500, detail="文件保存失败")

def validate_id_card(id_card: str) -> bool:
    """验证身份证号格式"""
    if not id_card:
        return True  # 允许为空
        
    # 身份证号格式正则表达式
    pattern = r'^[1-9]\d{5}(19|20)\d{2}(0[1-9]|1[0-2])(0[1-9]|[12]\d|3[01])\d{3}[\dXx]$'
    return bool(re.match(pattern, id_card))

@app.post("/faces/add", tags=["人脸管理"])
async def add_face(
    file: UploadFile = File(...),
    person_id: str = Form(...),
    name: Optional[str] = Form(None),
    age: Optional[str] = Form(None),
    gender: Optional[str] = Form(None),
    id_card: Optional[str] = Form(None),
    save_image: bool = Form(False)
):
    """
    添加人脸到数据库
    
    - **file**: 人脸图片文件
    - **person_id**: 人物ID
    - **name**: 人物姓名（可选）
    - **age**: 年龄（可选）
    - **gender**: 性别（可选）
    - **id_card**: 身份证号（可选）
    - **save_image**: 是否保存原始图片（可选，默认False）
    """
    try:
        # 验证身份证号格式
        if id_card and not validate_id_card(id_card):
            raise HTTPException(
                status_code=400,
                detail="身份证号格式无效"
            )
            
        # 检查文件是否为空
        if not file.filename:
            raise HTTPException(status_code=400, detail="未提供文件")
            
        # 保存文件
        file_path = save_upload_file(file, permanent=save_image)
        
        # 处理年龄字段
        age_value = None
        if age:
            try:
                age_value = int(age)
                if age_value < 0 or age_value > 150:  # 添加合理的年龄范围检查
                    raise ValueError("年龄必须在0-150之间")
            except ValueError as e:
                raise HTTPException(status_code=400, detail=f"年龄格式无效: {str(e)}")
        
        # 构建元数据
        metadata = {
            'name': name,
            'age': age_value,  # 使用转换后的年龄值
            'gender': gender,
            'id_card': id_card,
            'original_filename': file.filename,
            'saved_path': file_path if save_image else None
        }
        metadata = {k: v for k, v in metadata.items() if v is not None}
            
        # 添加人脸
        success = searcher.add_face(file_path, person_id, metadata)
        
        # 如果不是永久保存，则删除临时文件
        if not save_image:
            try:
                os.remove(file_path)
            except Exception as e:
                logger.warning(f"删除临时文件失败: {str(e)}")
        
        if success:
            return {
                "status": "success", 
                "message": "人脸添加成功",
                "saved_path": file_path if save_image else None
            }
        else:
            # 如果是永久保存但添加失败，也需要删除文件
            if save_image:
                try:
                    os.remove(file_path)
                except Exception as e:
                    logger.warning(f"删除文件失败: {str(e)}")
            raise HTTPException(
                status_code=400, 
                detail="未能检测到人脸或人脸添加失败"
            )
            
    except HTTPException:
        raise
    except Exception as e:
        # 确保清理临时文件
        if 'file_path' in locals():
            try:
                os.remove(file_path)
            except:
                pass
        logger.error(f"添加人脸失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"添加人脸失败: {str(e)}")

# 添加超时中间件
@app.middleware("http")
async def timeout_middleware(request: Request, call_next):
    try:
        return await asyncio.wait_for(call_next(request), timeout=60.0)  # 60秒超时
    except asyncio.TimeoutError:
        raise HTTPException(status_code=504, detail="请求超时")

@app.post("/faces/search", tags=["人脸搜索"])
async def search_faces(
    file: UploadFile = File(...),
    threshold: float = Form(0.6),
    top_k: int = Form(5)
):
    """
    搜索相似人脸
    
    - **file**: 查询用的人脸图片
    - **threshold**: 相似度阈值 (0-1)
    - **top_k**: 返回的最相似人脸数量
    """
    try:
        if not file.filename:
            raise HTTPException(status_code=400, detail="未提供文件")
            
        # 保存上传的文件
        file_path = save_upload_file(file)
        
        try:
            # 设置搜索超时
            results = await asyncio.wait_for(
                searcher.search_similar_faces_async(
                    file_path, 
                    threshold, 
                    top_k
                ),
                timeout=30.0  # 30秒超时
            )
            
            if results is None:
                raise HTTPException(
                    status_code=400,
                    detail="人脸库中未查询到当前人脸"
                )
            
            return {
                "status": "success",
                "results": results
            }
            
        except asyncio.TimeoutError:
            raise HTTPException(
                status_code=504,
                detail="搜索超时，请稍后重试"
            )
        finally:
            try:
                os.remove(file_path)
            except Exception as e:
                logger.warning(f"删除临时文件失败: {str(e)}")
                
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"搜索人脸失败: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"搜索人脸失败: {str(e)}"
        )

@app.post("/faces/batch-add", tags=["人脸管理"])
async def batch_add_faces(
    files: List[UploadFile] = File(...),
    person_id: str = Form(...),
    name: Optional[str] = Form(None),
    age: Optional[str] = Form(None),
    gender: Optional[str] = Form(None),
    id_card: Optional[str] = Form(None)
):
    """
    批量添加人脸
    
    - **files**: 人脸图片文件列表
    - **person_id**: 人物ID
    - **name**: 人物姓名（可选）
    - **age**: 年龄（可选）
    - **gender**: 性别（可选）
    - **id_card**: 身份证号（可选）
    """
    try:
        # 验证身份证号格式
        if id_card and not validate_id_card(id_card):
            raise HTTPException(
                status_code=400,
                detail="身份证号格式无效"
            )
            
        # 处理年龄字段
        age_value = None
        if age:
            try:
                age_value = int(age)
                if age_value < 0 or age_value > 150:
                    raise ValueError("年龄必须在0-150之间")
            except ValueError as e:
                raise HTTPException(status_code=400, detail=f"年龄格式无效: {str(e)}")
        
        metadata = {
            'name': name,
            'age': age_value,
            'gender': gender,
            'id_card': id_card
        }
        metadata = {k: v for k, v in metadata.items() if v is not None}
        
        success_count = 0
        failed_files = []
        
        for file in files:
            file_path = save_upload_file(file)
            try:
                if searcher.add_face(file_path, person_id, metadata):
                    success_count += 1
                else:
                    failed_files.append(file.filename)
            finally:
                os.remove(file_path)
                
        return {
            "status": "success",
            "success_count": success_count,
            "failed_files": failed_files
        }
        
    except Exception as e:
        logger.error(f"批量添加人脸失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/faces/{person_id}", tags=["人脸管理"])
async def delete_person(person_id: str):
    """
    删除指定人物的所有人脸记录
    
    - **person_id**: 要删除的人物ID
    """
    try:
        deleted_count = searcher.delete_person(person_id)
        return {
            "status": "success",
            "deleted_count": deleted_count
        }
    except Exception as e:
        logger.error(f"删除人脸失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health", tags=["系统"])
async def health_check():
    """健康检查接口"""
    try:
        # 检查 Milvus 连接
        client = searcher.get_client()
        client.list_collections()
        return {"status": "healthy"}
    except Exception as e:
        logger.error(f"健康检查失败: {str(e)}")
        raise HTTPException(
            status_code=503,
            detail="服务不可用"
        )

def signal_handler(signum, frame):
    """处理退出信号"""
    logger.info("接收到退出信号，正在清理资源...")
    try:
        if hasattr(searcher, 'cleanup'):
            searcher.cleanup()
        logger.info("资源清理完成，正在退出...")
    except Exception as e:
        logger.error(f"清理资源时出错: {str(e)}")
    finally:
        sys.exit(0)

# 注册信号处理器
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

def start_server(host="0.0.0.0", port=8000):
    """启动服务器"""
    try:
        logger.info("启动服务器...")
        uvicorn.run(app, host=host, port=port)
    except Exception as e:
        logger.error(f"服务器启动失败: {str(e)}")
    finally:
        logger.info("服务器关闭")
        if hasattr(searcher, 'cleanup'):
            searcher.cleanup()

if __name__ == "__main__":
    try:
        start_server()
    except KeyboardInterrupt:
        logger.info("接收到中断信号")
        sys.exit(0)
    except Exception as e:
        logger.error(f"程序异常退出: {str(e)}")
        sys.exit(1) 